Seismic search optimizations by chishui · Pull Request #20 · opensearch-project/neural-sparse-cpp

chishui · 2026-07-23T06:53:21Z

Description

Done some optimization which could bring overall query performance by >30%

k = 10

  ┌──────────────────────┬─────────────────────┬─────────┐
  │        Config        │ CPU ms (mean ± std) │ Speedup │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ baseline             │ 127.7 ± 12.4        │ 1.00×   │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ + VisitedSet only    │ 116.6 ± 1.5         │ 1.10×   │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ + head-prefetch only │ 107.6 ± 8.9         │ 1.19×   │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ + hugepage only      │ 124.8 ± 6.0         │ 1.02×   │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ all three together   │ 91.9 ± 0.3          │ 1.39×   │
  └──────────────────────┴─────────────────────┴─────────┘

k = 100

  ┌──────────────────────┬─────────────────────┬─────────┐
  │        Config        │ CPU ms (mean ± std) │ Speedup │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ baseline             │ 353.6 ± 8.8         │ 1.00×   │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ + VisitedSet only    │ 347.4 ± 6.1         │ 1.02×   │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ + head-prefetch only │ 291.2 ± 5.6         │ 1.21×   │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ + hugepage only      │ 336.9 ± 3.9         │ 1.05×   │
  ├──────────────────────┼─────────────────────┼─────────┤
  │ all three together   │ 257.8 ± 1.7         │ 1.37×   │
  └──────────────────────┴─────────────────────┴─────────┘

Notes

head-prefetch is the dominant lever at both k (1.19× / 1.21×).
VisitedSet helps meaningfully at k=10 (1.10×) but washes out at k=100 (1.02×, within noise).
hugepage alone is small (1.02× / 1.05×) — borderline noise.
All three together (1.39× / 1.37×) exceed the levers stacked individually — they're synergistic: once the doc-scan is prefetch-bound, the dedup and page-walk wins compound rather than add.
"All three" is the exact config now committed (d704f61); the two dropped levers (doc-offset precompute, two-lane AVX512) were within noise and removing them cost only ~1%.

Issues Resolved

List any issues this PR will resolve, e.g. Closes [...].

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

yuye-aws · 2026-07-23T08:40:49Z

+// smaller (1 bit/doc, ~1.1 MB at 8.8M docs) so far fewer distinct cache lines
+// are touched, while the touched-word list lets a new query clear only the
+// words it actually dirtied (sparse O(visited) reset, not O(n)).
+class VisitedSet {


Thanks for the detailed estimation on the 8.8M dataset. Do you think the current visited set can apply to all dataset scales? I'm worrying the bit set size from the super large dataset

It scales with data size, so 11 MB for 88M, and we've benchmark at most 45M docs for a single host, seems OK to me, but I'll benchmark on 45M later.

I understand users rarely use more than 45M in a single segment, but there is no guarantee that the segment size will always be smaller than a certain threshold. How about also keeping the hash set solution and switch the set solution based on the dataset size? For example, dataset size <= 100M will use the bit set.

When we are not even testing a larger dataset, why should we set a presumable hard coded number to branch two logic.

We're not testing a larger dataset, but anything could happen from the user side where they perform force merge operation so that the segment contains soo many documents that the visited set consumes too much memory.

I also don't feel like the hard coded number solution, but memory issue always exists. Even if the C++ solution is working on non-JVM memory. The users can still limit their circuit breaker setting, making it not possible to use the bitset solution.

zirui-song-18

In addition, I'm worried that the huge pages might introduce P99 spikes via NUMA auto-migration. Would you mind testing the P50 P90 P99 data before & after the optimization?

zirui-song-18 · 2026-07-23T08:38:02Z

+    void new_query() {
+        for (const size_t w : touched_) {
+            bits_[w] = 0;
+        }
+        touched_.clear();
+    }


The new_query() sparse-reset is the most fragile point in this PR — it clears only touched_ words, relying on "a word is in touched_ iff it has a set bit" (maintained solely by if (prev==0) touched_.push_back(w)). I don't see a VisitedSet test (the only test change is transposed_scoring_matches_reference). Please add: (a) dedup within a query; (b) full reset across new_query; (c) two ids in the same 64-bit word push touched_ once and both clear; (d) resize() clears touched_; (e) boundaries 63/64/n-1. A regression here silently corrupts recall and is invisible to perf benchmarks.

yuye-aws · 2026-07-23T08:43:45Z

+    void new_query() {
+        for (const size_t w : touched_) {
+            bits_[w] = 0;
+        }
+        touched_.clear();
+    }


When k is larger than 100, new_query will take longer time clearing the touched. Can we have different strategy for the set? For example, when k <= 100, use the current one. If k > 1000, use the hash set?

It's actually related to the doc count we visited, not the value of k, right?

Larger k will visit more documents

Oh, you are right about k visiting more documents, but the iteration of touched is quite cache friendly by sequentially traverse, right? I don't think it makes quite sense to set a hard-coded number k to introduce branching logic make less readible code and unpredictable branches.

For k = 10, VisitedSet only boost is 1.10×. For k = 100, VisitedSet only boost is 1.02×. I am afraid that for a larger value of k, the visited set boost will have negative impact.

Inter-process ablation on the 8.8M-doc corpus shows ~1.4x for float Seismic and ~1.06x for the 8-bit SQ variant from: - VisitedSet: bitset dedup with sparse per-query reset, replacing the absl::flat_hash_set on the doc-scan critical path. - prefetch_vector_head: prefetch only the leading cache lines of the next doc row, bounding outstanding software prefetches. - MADV_HUGEPAGE/COLLAPSE the corpus arrays to cut TLB/page-walk cost on the random per-doc gather. - sort_cluster_docs: sort doc ids ascending within each cluster so the per-doc gather is monotonic. Doc-offset precompute and two-lane AVX512 were within noise, so omitted. Add VisitedSet unit tests for the fragile sparse-reset: within-query dedup, full reset across new_query, two ids sharing a 64-bit word, resize clearing touched state, and word boundaries (63/64/n-1). Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

VisitedSet indexes by doc id with no per-candidate bounds check (a branch there would erode the dedup win), so "every doc id < num_vectors" is now a load-bearing invariant: an out-of-range id is an out-of-bounds write, not the safe insert it was with the former hash set. Add InvertedListClusters::validate_doc_ids(num_docs), an O(nnz) scan that throws std::out_of_range on a stray id, called from both SeismicIndex and SeismicScalarQuantizedIndex read_index() after load. A debug assert in VisitedSet::insert documents the invariant at the point of use. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

chishui · 2026-07-28T05:12:15Z

This is the performance comparison on 46M MS Marco V2 (1/3 corpus size)

 Single-thread query benchmark: opt vs main, full 46M, scaled params

  seismic-ec2-xl (96c/184GB), 46,121,400-doc scaled index (λ=31298, β=2087), msmarco-v2-queries.csr (3903 queries, dim-matched = 30522), OMP_NUM_THREADS=1, 10 reps:

  ┌─────┬──────────────────────────┬─────────────────────────┬───────────────┐
  │  k  │ main QPS (mean / median) │ opt QPS (mean / median) │    Speedup    │
  ├─────┼──────────────────────────┼─────────────────────────┼───────────────┤
  │ 10  │ 145.9 / 156.0            │ 224.7 / 246.6           │ 1.54× / 1.58× │
  ├─────┼──────────────────────────┼─────────────────────────┼───────────────┤
  │ 100 │ 72.3 / 69.6              │ 110.8 / 123.3           │ 1.53× / 1.77× │
  └─────┴──────────────────────────┴─────────────────────────┴───────────────┘

GCC in the Linux CI container does not pull global ::size_t in via <cstdint>/<vector>, so every unqualified size_t use failed to compile. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

chishui · 2026-07-28T08:25:06Z

In addition, I'm worried that the huge pages might introduce P99 spikes via NUMA auto-migration. Would you mind testing the P50 P90 P99 data before & after the optimization?

┌──────┬────────────────────────┬────────────────────────┬────────────────────────┬─────────────────────────┐
  │  k   │          p50           │          p90           │          p99           │           max           │
  ├──────┼────────────────────────┼────────────────────────┼────────────────────────┼─────────────────────────┤
  │ 10   │ 0.70 → 0.65 (−7.7%)    │ 1.85 → 1.57 (−14.8%)   │ 3.72 → 2.78 (−25.4%)   │ 9.28 → 4.90 (−47.2%)    │
  ├──────┼────────────────────────┼────────────────────────┼────────────────────────┼─────────────────────────┤
  │ 100  │ 2.45 → 1.99 (−18.7%)   │ 5.21 → 4.14 (−20.5%)   │ 8.44 → 6.75 (−20.0%)   │ 16.34 → 13.19 (−19.3%)  │
  ├──────┼────────────────────────┼────────────────────────┼────────────────────────┼─────────────────────────┤
  │ 1000 │ 9.34 → 7.47 (−20.0%)   │ 15.75 → 12.29 (−22.0%) │ 23.36 → 16.18 (−30.7%) │ 79.27 → 23.25 (−70.7%)  │
  ├──────┼────────────────────────┼────────────────────────┼────────────────────────┼─────────────────────────┤
  │ 5000 │ 18.39 → 16.02 (−12.9%) │ 27.39 → 22.57 (−17.6%) │ 45.36 → 27.07 (−40.3%) │ 309.75 → 61.25 (−80.2%) │
  └──────┴────────────────────────┴────────────────────────┴────────────────────────┴─────────────────────────┘

chishui · 2026-07-28T09:17:46Z

Head-prefetch seems to be the only consistent optimization

 ┌──────────────────────┬──────────────────────┬──────────────────────┬──────────────────────┐
  │        Lever         │         k=10         │        k=100         │        k=5000        │
  ├──────────────────────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ + VisitedSet only    │ 0.98× (t=−1.6) noise │ 1.00× (t=−0.4) noise │ 1.05× (t=+1.5) noise │
  ├──────────────────────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ + head-prefetch only │ 1.15× (t=+10.7)      │ 1.19× (t=+21.2)      │ 1.15× (t=+4.6)       │
  ├──────────────────────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ + hugepage only      │ 1.01× (t=+0.9) noise │ 1.00× (t=+0.1) noise │ 0.97× (t=−0.8) noise │
  ├──────────────────────┼──────────────────────┼──────────────────────┼──────────────────────┤
  │ all together         │ 1.20×                │ 1.21×                │ 1.26× (t=+7.5)       │
  └──────────────────────┴──────────────────────┴──────────────────────┴──────────────────────┘

But VisitedSet does earn its place — in the tail, not the mean. At k=5000 vs baseline, we could do thorough experiment with it later.

  ┌────────────────────┬──────┬──────┬──────┬──────┐
  │       Lever        │ p50  │ p90  │ p99  │ max  │
  ├────────────────────┼──────┼──────┼──────┼──────┤
  │ VisitedSet only    │ +1%  │ −2%  │ −21% │ −79% │
  ├────────────────────┼──────┼──────┼──────┼──────┤
  │ head-prefetch only │ −22% │ −23% │ −14% │ −7%  │
  ├────────────────────┼──────┼──────┼──────┼──────┤
  │ all together       │ −15% │ −18% │ −39% │ −79% │
  └────────────────────┴──────┴──────┴──────┴──────┘

will only keep head-prefetch

Per-lever ablation on base_full (8.8M docs, 72 threads, 20 reps, Welch t-test on batch times) showed head-prefetch is the only lever with a statistically significant effect: lever k=10 k=100 head-prefetch 1.15x (t=+10.7) 1.19x (t=+21.2) VisitedSet 0.98x (t=-1.6) 1.00x (t=-0.4) hugepage 1.01x (t=+0.9) 1.00x (t=+0.1) per-cluster sort 0.98x (t=-1.6) 0.98x (t=-2.3) scratch hoisting 0.97x (t=-2.5) 1.01x (t=+0.5) All levers together gave 1.20x/1.21x, i.e. no synergy beyond prefetch alone. Drop the levers that do not pay for their complexity: the bitset VisitedSet (and its load-time doc-id validation, needed only because it indexed without bounds checks), the hugepage madvise, the per-cluster doc-id sort, and the per-thread scratch hoisting. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

chishui requested review from model-collapse, yuye-aws and zirui-song-18 as code owners July 23, 2026 06:53

chishui force-pushed the seismic-search-optimizations-and-mmap branch from 17417f6 to ce6a599 Compare July 23, 2026 07:00

yuye-aws reviewed Jul 23, 2026

View reviewed changes

zirui-song-18 reviewed Jul 23, 2026

View reviewed changes

yuye-aws reviewed Jul 23, 2026

View reviewed changes

chishui added 3 commits July 23, 2026 09:17

Raise gtest discovery timeout to avoid flaky Windows CI failures

7037978

Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

chishui force-pushed the seismic-search-optimizations-and-mmap branch from 0125691 to 7037978 Compare July 28, 2026 05:19

Include <cstddef> for size_t in visited_set.h

5614de5

GCC in the Linux CI container does not pull global ::size_t in via <cstdint>/<vector>, so every unqualified size_t use failed to compile. Signed-off-by: Liyun Xiu <xiliyun@amazon.com>

yuye-aws approved these changes Jul 30, 2026

View reviewed changes

yuye-aws merged commit 85a7be7 into opensearch-project:main Jul 30, 2026
8 checks passed

Uh oh!

Conversation

chishui commented Jul 23, 2026

Description

Issues Resolved

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zirui-song-18 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

chishui commented Jul 28, 2026

Uh oh!

chishui commented Jul 28, 2026

Uh oh!

chishui commented Jul 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants